import os
import random

import numpy as np
from PIL import Image
from torchvision.transforms import transforms

from seg_dataset import RandomHorizontalFlip, RandomRotation

img_h,img_w=480,480
imgs=np.zeros([3,img_h,img_w,1])
means,std=[0,0,0],[0,0,0]
R = random.uniform(0,1)
path=r'D:\study\pytorch_study\seg_thryoid_picture\datasets\image'
data_transform=transforms.Compose([

            transforms.Resize((480,480)),
            RandomHorizontalFlip(0.5,R),
            RandomRotation(0.5,R),
            transforms.CenterCrop(480),
            transforms.ToTensor(),

        ])
print(len(os.listdir(path)))
i=0
for img_name in os.listdir(path):
    # print(img_name)
    i+=1
    try:
        img=Image.open(os.path.join(path,img_name)).convert('RGB')
        img=data_transform(img)
        img=img.numpy()

        # img=img[:,:,:,np.newaxis]
        # imgs=np.concatenate((imgs,img),axis=3)
    except(OSError,NameError):
        print('OSError')
    for i in range(3):
        means[i] += img[i, :, :].mean()
        std[i] += img[i, :, :].std()
    if i%1000==0:
        print(str(i))
means = np.asarray(means) /len(os.listdir(path))
std = np.asarray(std) / len(os.listdir(path))
print(means)
print(std)